Update app.py

app.py

@@ -17,6 +17,8 @@ from PIL import ImageDraw, UnidentifiedImageError
 from PIL import Image as PILImage
 
 
+image_dir = "saved_images"
+os.makedirs(image_dir, exist_ok=True)
 base_url = "https://huggingface.co/spaces/Tonic1/Official-Qwen-VL-Chat"
 model_name = "Qwen/Qwen-VL-Chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
@@ -43,21 +45,27 @@ class ChatBot:
         response, self.history = self.model.chat(self.tokenizer, query=query, history=self.history)
         return response
 
-    def draw_boxes(self, response):
-        image = self.tokenizer.draw_bbox_on_latest_picture(response, self.history)
-        if image is not None:
-            if not isinstance(image, PILImage.Image):
-                image = PILImage.fromarray(image.get_image())
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            return "data:image/png;base64," + img_str
-        else:
+    def draw_boxes(self, response, image_path):
+        boxes = re.findall(r'<box>\((\d+),(\d+)\),\((\d+),(\d+)\)</box>', response)
+        if not boxes:
             return None
-
+        try:
+            with PILImage.open(image_path) as img:
+                draw = ImageDraw.Draw(img)
+                for box in boxes:
+                    x1, y1, x2, y2 = map(int, box)
+                    draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+                file_name = secrets.token_hex(10) + ".png"
+                file_path = os.path.join(image_dir, file_name)
+                img.save(file_path, format="PNG")
+                return file_path
+        except Exception as e:
+            print(f"An error occurred while processing the image: {e}")
+            return None
+            
     def clean_response(self, response):
         return re.sub(r'<ref>(.*?)</ref>(?:<box>.*?</box>)*(?:<quad>.*?</quad>)*', r'\1', response).strip()
-
+    
     def clear_memory(self):
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -69,10 +77,10 @@ def chat_interface(text_query, file):
     response = chatbot.chat(image_path=image_path, text_query=text_query)
 
     if "<box>" in response:
-        image_with_boxes = chatbot.draw_boxes(response)
+        image_file_path = chatbot.draw_boxes(response, image_path)
         text_response = chatbot.clean_response(response)
         chatbot.clear_memory()
-        return [("Qwen-VL_Chat", text_response), ("Qwen-VL_Image", image_with_boxes)]
+        return [("Qwen-VL_Chat", text_response), ("Qwen-VL_Image", image_file_path)]
     else:
         chatbot.clear_memory()
         return [("Qwen-VL_Chat", response)]
@@ -108,4 +116,4 @@ Note: This demo is governed by the original license of Qwen-VL. We strongly advi
 including hate speech, violence, pornography, deception, etc. (Note: This demo is subject to the license agreement of Qwen-VL. We strongly advise users not to disseminate or allow others to disseminate the following content, including but not limited to hate speech, violence, pornography, and fraud-related harmful information.)
 """)
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()