add ration and site parameters

imagecrawler.go

@@ -11,8 +11,10 @@ import (
   "net/url"
   "os"
   "os/user"
+  "strconv"
   "strings"
   "sync"
+  "time"
 )
 
 // Global vars
@@ -20,6 +22,12 @@ var wg sync.WaitGroup
 
 type picture struct {
   FileURL string `json:"file_url"`
+  Width        int    `json:"width"`
+  Height       int    `json:"height"`
+  ImageWidth   int    `json:"image_width"`
+  ImageHeight  int    `json:"image_height"`
+  FileExt      string `json:"file_ext"`
+  Tags         string `json:"tag_string"`
 }
 
 // main function to download pictures
@@ -29,14 +37,19 @@ func main() {
   var path string
   var safemode bool
   var tags string
+  var aspect string
+  var site string
 
   // variables for downloading
   picHits := 1
   page := 1
+  ratio := 0.0
 
   flag.StringVar(&path, "dir", "unnamed", "Directory to safe pictures. Default is %HOME/pictures/konachan/unnamed")
   flag.BoolVar(&safemode, "safe", false, "Safemode to filter NSFW pictures. Default is false")
   flag.StringVar(&tags, "tags", "", "Tags used to filter search query.")
+  flag.StringVar(&aspect, "aspect", "", "Aspect ratio pics should have")
+  flag.StringVar(&site, "site", "konachan", "Site to crawl from, either konachan or danbooru")
 
   flag.Parse()
 
@@ -45,7 +58,7 @@ func main() {
   if err != nil {
     log.Fatal(err)
   }
-	filepath := strings.Join([]string{homepath.HomeDir, "pictures", "konachan", strings.TrimSuffix(path, "\n")}, "/")
+  filepath := strings.Join([]string{homepath.HomeDir, "pictures", site, strings.TrimSuffix(path, "\n")}, "/")
   os.MkdirAll(filepath, 0700)
 
   // edit tags array to met API requirement
@@ -53,6 +66,16 @@ func main() {
   tags = strings.Replace(tags, "=", ":", -1)
   tags = strings.TrimSuffix(tags, "\n")
 
+  // calculate aspect ratio
+  if isFlagPassed("aspect") {
+    aspectSlice := strings.Split(aspect, ":")
+    widthF, _ := strconv.ParseFloat(aspectSlice[0], 64)
+    heightF, _ := strconv.ParseFloat(aspectSlice[1], 64)
+    ratio = widthF / heightF
+  } else {
+    ratio = 0.0
+  }
+
   for picHits > 0 {
 
     fmt.Println("Page: ", page)
@@ -62,10 +85,17 @@ func main() {
       website = fmt.Sprintf("https://konachan.com/post.json?page=%d&tags=%s+rating:safe", page, tags)
     }
 
-		picList := openConnection(website)
-		pictures := parseMaps(picList)
+    if site == "danbooru" {
+      website = fmt.Sprintf("https://danbooru.donmai.us/posts.json?page=%d&tags=%s", page, tags)
+      if safemode {
+        website = fmt.Sprintf("https://danbooru.donmai.us/posts.json?page=%d&tags=%s+rating:safe", page, tags)
+      }
+    }
 
-		picHits = len(pictures)
+    picList := openConnection(website)
+    pictures, count := parseMaps(picList, ratio)
+
+    picHits = count
     page++
 
     wg.Add(len(pictures))
@@ -73,9 +103,20 @@ func main() {
       go downloadPic(pic, filepath)
     }
     wg.Wait()
+    time.Sleep(1 * time.Second)
   }
 }
 
+func isFlagPassed(name string) bool {
+  found := false
+  flag.Visit(func(f *flag.Flag) {
+    if f.Name == name {
+      found = true
+    }
+  })
+  return found
+}
+
 // function to create the connection to konachan and get the API response
 func openConnection(url string) []picture {
   var f []picture
@@ -99,14 +140,38 @@ func openConnection(url string) []picture {
 }
 
 // function to parse the json response and extract only the file url
-func parseMaps(f []picture) []string {
+func parseMaps(f []picture, ratio float64) ([]string, int) {
   fileURLs := []string{}
+  picCount := 0
+
+  if isFlagPassed("aspect") {
     for _, pic := range f {
+      picCount++
+
+      picWidthF := 1.0
+      picHeightF := 1.0
+
+      if (pic.Width != 0 && pic.Height != 0) {
+        picWidthF = float64(pic.Width)
+        picHeightF = float64(pic.Height)
+      } else {
+        picWidthF = float64(pic.ImageWidth)
+        picHeightF = float64(pic.ImageHeight)
+      }
+
+      if (picWidthF / picHeightF) == ratio {
         fileURL := pic.FileURL
         fileURLs = append(fileURLs, fileURL)
       }
-
-	return fileURLs
+    }
+  } else {
+    for _, pic := range f {
+      picCount++
+      fileURL := pic.FileURL
+      fileURLs = append(fileURLs, fileURL)
+    }
+  }
+  return fileURLs, picCount
 }
 
 // function to download and sace the pictures to disk