Update app.py

app.py

@@ -91,8 +91,7 @@ example_images = [
 ]
 
 
-# Main function to process the uploaded image
-def process_image(img, generate_image=True):
+def process_classification(img):
     print("Starting prediction...")
     predicted_class, _, probs = learn.predict(img)
     print(f"Prediction complete: {predicted_class}")
@@ -109,33 +108,26 @@ def process_image(img, generate_image=True):
     endangerment_status = get_status(predicted_class)
     print(f"Status found: {endangerment_status}")
 
-    # Generate artistic interpretation using DALL-E
+    return classification_results, wiki_url, endangerment_status, predicted_class
+
+def generate_artistic_image(predicted_class):
     print("Sending request to DALL-E...")
     try:
         client = OpenAI()
+        result = client.images.generate(
+            model="gpt-image-1",
+            prompt=random.choice(prompt_templates).format(flower=predicted_class),
+            size="1024x1024"
+        )
 
-        if generate_image:
-            result = client.images.generate(
-                model="gpt-image-1",
-                prompt=random.choice(prompt_templates).format(flower=predicted_class),
-                size="1024x1024",
-                background="transparent",
-                quality="low"
-            )
-
-            image_base64 = result.data[0].b64_json
-            image_bytes = base64.b64decode(image_base64)
-            generated_image = Image.open(io.BytesIO(image_bytes))
-        else:
-            generated_image = None
+        image_base64 = result.data[0].b64_json
+        image_bytes = base64.b64decode(image_base64)
+        generated_image = Image.open(io.BytesIO(image_bytes))
+        return generated_image
 
     except Exception as e:
         print(f"Error generating image: {e}")
-        generated_image = None
-
-    print("Image retrieved and ready to return")
-    return classification_results, generated_image, wiki_url, endangerment_status
-
+        return None
 
 # Function to clear all outputs
 def clear_outputs():
@@ -173,10 +165,24 @@ with gr.Blocks() as demo:
         outputs=None
     )
 
+    # Store the predicted class for image generation
+    predicted_class = gr.State()
+
+    def process_and_generate(img):
+        # First get classification results
+        classification_results, wiki_url, endangerment_status, pred_class = process_classification(img)
+        
+        # Return classification results immediately
+        yield classification_results, None, wiki_url, endangerment_status, pred_class
+        
+        # Then generate and return the image
+        generated_img = generate_artistic_image(pred_class)
+        yield classification_results, generated_img, wiki_url, endangerment_status, pred_class
+
     input_image.change(
-        fn=lambda img: process_image(img, generate_image=True),
+        fn=process_and_generate,
         inputs=input_image,
-        outputs=[label_output, generated_image, wiki_output, endangerment_output]
+        outputs=[label_output, generated_image, wiki_output, endangerment_output, predicted_class]
     )
 
     input_image.clear(