sending image and json with captions

api/app.py

@@ -5,6 +5,7 @@ import os
 import io
 from model.src.utils.arg_parser import eval_parse_args  # Nuovo import corretto
 import sys
+from PIL import Image
 
 from model.src import eval
 
@@ -15,17 +16,34 @@ CORS(app)
 def index():
     return render_template('index.html')
 
-@app.route('/generate-design', methods=['POST'])
+@app.route('/generate-design', methods=['GET','POST'])
 def generate_design():
     try:
+        
+        # Getting json
+        json_data_from_req = request.get_json()
+        if not json_data_from_req:
+            return "Invalid or missing JSON data", 400
+        print(json_data_from_req)
+
 
         # Getting Image
-        image_file = request.files['image']
-        image = Image.open(image_file)
-        save_path = os.path.join('/api/model/assets/data/vitonhd/test/im_sketch', '03191_00.jpeg')
-        image.save(save_path, 'JPEG')
 
+        if 'image' not in request.files:
+            return "No image file in request", 400
 
+        image_file = request.files['image']
+        try:
+            image = Image.open(image_file)
+        except Exception as e:
+            return f"Failed to open the image: {str(e)}", 400
+        
+        # Create an in-memory buffer to store the image (instead of saving to disk)
+        img_sketch_buffer = io.BytesIO()
+        # Save the image to the buffer in JPEG format
+        image.save(img_sketch_buffer, format='JPEG')
+        # Rewind the buffer's position to the beginning
+        img_sketch_buffer.seek(0)
 
         
         # Creiamo una lista di argomenti come quelli che passeresti via CLI
@@ -42,7 +60,7 @@ def generate_design():
         ]
 
          # Esegui la funzione `main()` di eval.py passando gli argomenti
-        final_image = eval.main()
+        final_image = eval.main(img_sketch_buffer, json_data_from_req)
 
         # Save the image to a BytesIO buffer to return via HTTP
         img_io = io.BytesIO()

api/model/src/datasets/vitonhd.py

@@ -26,6 +26,8 @@ class VitonHDDataset(data.Dataset):
             self,
             dataroot_path: str,
             phase: str,
+            im_sketch_buffer_from_request,
+            json_from_req,
             tokenizer,
             radius=5,
             caption_folder='captions.json', #######################################################3
@@ -48,6 +50,8 @@ class VitonHDDataset(data.Dataset):
         self.width = size[1]
         self.radius = radius
         self.tokenizer = tokenizer
+        self.im_sketch_buffer_from_request = im_sketch_buffer_from_request
+        self.json_from_req = json_from_req
         self.transform = transforms.Compose([
             transforms.ToTensor(),
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
@@ -70,12 +74,19 @@ class VitonHDDataset(data.Dataset):
         assert all(x in possible_outputs for x in outputlist)
 
         # Load Captions
+        model_data  = self.json_from_req.get('MODEL', {})  # Safely get the 'MODEL' key, default to an empty dictionary if it doesn't exist
+        # Filter captions based on the length requirement (3 or more items)
+        self.captions_dict = {k: v for k, v in model_data.items() if len(v) >= 3}
+        '''
         with open(os.path.join(self.dataroot, self.caption_folder)) as f:
             # self.captions_dict = json.load(f)['items']
             self.captions_dict = json.load(f)
         self.captions_dict = {k: v for k, v in self.captions_dict.items() if len(v) >= 3}
+        
 
         dataroot = self.dataroot
+
+        
         if phase == 'train':
             filename = os.path.join(dataroot, f"{phase}_pairs.txt")
         else:
@@ -99,10 +110,10 @@ class VitonHDDataset(data.Dataset):
                 im_names.append(im_name)
                 c_names.append(c_name)
                 dataroot_names.append(dataroot)
-
-        self.im_names = im_names
-        self.c_names = c_names
-        self.dataroot_names = dataroot_names
+        '''
+        self.im_names = []
+        self.c_names = []
+        self.dataroot_names = []
 
     def __getitem__(self, index):
         """
@@ -112,9 +123,10 @@ class VitonHDDataset(data.Dataset):
         :return: dict containing dataset samples
         :rtype: dict
         """
-        c_name = self.c_names[index]
-        im_name = self.im_names[index]
-        dataroot = self.dataroot_names[index]
+        c_name = list(self.captions_dict.keys())[0] + "_00.jpg" # self.c_names[index]
+        im_name =  list(self.captions_dict.keys())[0] + "_00.jpg" #self.im_names[index]
+        #dataroot = self.dataroot_names[index]
+        dataroot = "./assets/data/vitonhd"
 
         sketch_threshold = random.randint(self.sketch_threshold_range[0], self.sketch_threshold_range[1])
 
@@ -146,7 +158,7 @@ class VitonHDDataset(data.Dataset):
             image = self.transform(image)  # [-1,1]
 
         if "im_sketch" in self.outputlist:
-            # Person image
+            '''# Person image
             # im_sketch = Image.open(os.path.join(dataroot, 'im_sketch', c_name.replace(".jpg", ".png")))
             if self.order == 'unpaired':
                 im_sketch = Image.open(
@@ -161,8 +173,12 @@ class VitonHDDataset(data.Dataset):
             else:
                 raise ValueError(
                     f"Order should be either paired or unpaired"
-                )
-
+                )'''
+            im_sketch = Image.open(self.im_sketch_buffer_from_request)
+            # define a transform to convert the image to grayscale
+            transform = transforms.Grayscale()
+            # apply the above transform on the image
+            im_sketch = transform(im_sketch)
             im_sketch = im_sketch.resize((self.width, self.height))
             im_sketch = ImageOps.invert(im_sketch)
             # threshold grayscale pil image

api/model/src/eval.py

@@ -29,7 +29,7 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true"
 os.environ["WANDB_START_METHOD"] = "thread"
 
 
-def main() -> None:
+def main(im_sketch: io.BytesIO, json_data_from_req: Dict) -> None:
     args = eval_parse_args()
     accelerator = Accelerator(
         mixed_precision=args.mixed_precision,
@@ -91,6 +91,8 @@ def main() -> None:
             radius=5,
             tokenizer=tokenizer,
             size=(512, 384),
+            im_sketch=im_sketch,
+            json_data_from_req=json_data_from_req
         )
     else:
         raise NotImplementedError