Merge pull request #61 from Photoroom/ben/cmd_fix

[nit] Demo was stuck to 2000 samples..

cmd/main.go

@@ -11,7 +11,18 @@ import (
 )
 
 func main() {
-	// Define flags
+
+	cropAndResize := flag.Bool("crop_and_resize", false, "Whether to crop and resize the images and masks")
+	concurrency := flag.Int("concurrency", 64, "The number of concurrent http requests to make")
+	itemFetchBuffer := flag.Int("item_fetch_buffer", 256, "The number of items to pre-load")
+	itemReadyBuffer := flag.Int("item_ready_buffer", 128, "The number of items ready to be served")
+	limit := flag.Int("limit", 2000, "The number of items to fetch")
+	profile := flag.Bool("profile", false, "Whether to profile the code")
+
+	// Parse the flags before setting the configuration values
+	flag.Parse()
+
+	// Initialize the configuration
 	config := datago.GetDatagoConfig()
 
 	sourceConfig := datago.SourceFileSystemConfig{RootPath: os.Getenv("DATAGO_TEST_FILESYSTEM")}
@@ -22,18 +33,14 @@ func main() {
 	config.ImageConfig = datago.ImageTransformConfig{
 		DefaultImageSize:  1024,
 		DownsamplingRatio: 32,
-		CropAndResize:     *flag.Bool("crop_and_resize", false, "Whether to crop and resize the images and masks"),
+		CropAndResize:     *cropAndResize,
 	}
 	config.SourceConfig = sourceConfig
-	config.Concurrency = *flag.Int("concurrency", 64, "The number of concurrent http requests to make")
-	config.PrefetchBufferSize = *flag.Int("item_fetch_buffer", 256, "The number of items to pre-load")
-	config.SamplesBufferSize = *flag.Int("item_ready_buffer", 128, "The number of items ready to be served")
-	config.Limit = *flag.Int("limit", 2000, "The number of items to fetch")
-
-	profile := flag.Bool("profile", false, "Whether to profile the code")
+	config.Concurrency = *concurrency
+	config.PrefetchBufferSize = *itemFetchBuffer
+	config.SamplesBufferSize = *itemReadyBuffer
+	config.Limit = *limit
 
-	// Parse the flags and instantiate the client
-	flag.Parse()
 	dataroom_client := datago.GetClient(config)
 
 	// Go-routine which will feed the sample data to the workers

pkg/client.go

@@ -179,7 +179,7 @@ func GetClient(config DatagoConfig) *DatagoClient {
 
 	switch config.SourceConfig.(type) {
 	case SourceDBConfig:
-		fmt.Println("Creating a DB-backed dataloader")
+		fmt.Println("Creating a DB-backed dataloader.", config.Limit, " max samples")
 
 		var err error
 		dbConfig := config.SourceConfig.(SourceDBConfig)
@@ -190,7 +190,7 @@ func GetClient(config DatagoConfig) *DatagoClient {
 			backend = BackendHTTP{config: &dbConfig, concurrency: config.Concurrency}
 		}
 	case SourceFileSystemConfig:
-		fmt.Println("Creating a FileSystem-backed dataloader")
+		fmt.Println("Creating a FileSystem-backed dataloader.", config.Limit, " max samples")
 		fsConfig := config.SourceConfig.(SourceFileSystemConfig)
 		generator = newDatagoGeneratorFileSystem(fsConfig)
 		backend = BackendFileSystem{config: &config, concurrency: config.Concurrency}