Update app.py

app.py

@@ -91,7 +91,8 @@ example_images = [
 ]
 
 
-def process_classification(img):
+# Main function to process the uploaded image
+def process_image(img, generate_image=True):
     print("Starting prediction...")
     predicted_class, _, probs = learn.predict(img)
     print(f"Prediction complete: {predicted_class}")
@@ -108,26 +109,31 @@ def process_classification(img):
     endangerment_status = get_status(predicted_class)
     print(f"Status found: {endangerment_status}")
 
-    return classification_results, wiki_url, endangerment_status, predicted_class
-
-def generate_artistic_image(predicted_class):
+    # Generate artistic interpretation using DALL-E
     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"
-        )
 
-        image_base64 = result.data[0].b64_json
-        image_bytes = base64.b64decode(image_base64)
-        generated_image = Image.open(io.BytesIO(image_bytes))
-        return generated_image
+        if generate_image:
+            result = client.images.generate(
+                model="gpt-image-1",   
+                prompt=random.choice(prompt_templates).format(flower=predicted_class),
+                size="1024x1024"
+            )
+
+            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
 
     except Exception as e:
         print(f"Error generating image: {e}")
-        return None
+        generated_image = None
+
+    print("Image retrieved and ready to return")
+    return classification_results, generated_image, wiki_url, endangerment_status
+
 
 # Function to clear all outputs
 def clear_outputs():
@@ -165,24 +171,10 @@ 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=process_and_generate,
+        fn=lambda img: process_image(img, generate_image=True),
         inputs=input_image,
-        outputs=[label_output, generated_image, wiki_output, endangerment_output, predicted_class]
+        outputs=[label_output, generated_image, wiki_output, endangerment_output]
     )
 
     input_image.clear(