Update app.py

app.py

@@ -12,11 +12,13 @@ import zipfile
 
 # Set API tokens
 os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
-os.environ["OPENAI_API_KEY"] = "sk-proj-5iy4bwrqAW8GpguiEawaT3BlbkFJ8p88lLSjOCeDbxWsAOlr"
+os.environ["OPENAI_API_KEY"] = "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"
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Initialize the Replicate client
 rep_client = replicate.Client()
 
+# Predefined prompts for the dropdown
 predefined_prompts = [
     "Missing bolts on railway track",
     "Cracks on railway track",
@@ -26,11 +28,21 @@ predefined_prompts = [
     "Damaged railway platform"
 ]
 
-def handle_feedback(feedback, sentiment):
-    if sentiment == "like":
-        return "Thank you for your positive feedback!"
-    else:
-        return "Sorry to hear that. We are trying to improve based on your feedback."
+def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ'):
+    response = openai.ChatCompletion.create(
+        model=model_name,
+        messages=[
+            {
+                "role": "system",
+                "content": "The assistant is knowledgeable about rail defects and can answer questions related to them.",
+            },
+            {
+                "role": "user",
+                "content": question,
+            }
+        ],
+    )
+    return response.choices[0].message['content']
 
 def generate_variations(base_prompt, number_of_variations):
     locations = ["on the left side", "on the right side", "at the top", "at the bottom", "in the center"]
@@ -43,31 +55,19 @@ def generate_variations(base_prompt, number_of_variations):
         size = random.choice(sizes)
         weather = random.choice(weather_conditions)
         
-        full_prompt = f"{base_prompt}, with a {size} defect {location}, observed {weather}."
-        variations.append(full_prompt)
+        enhanced_prompt = f"{base_prompt}, with a {size} defect {location}, observed {weather}."
+        variations.append(enhanced_prompt)
     return variations
 
-def generate_images(prompts):
-    images = []
-    for prompt in prompts:
-        try:
-            prediction = rep_client.predictions.create(
-                version="ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e4",
-                input={"prompt": prompt, "scheduler": "K_EULER"}
-            )
-            prediction.wait()
-            if prediction.status == "succeeded" and prediction.output:
-                images.append(prediction.output[0])
-            else:
-                images.append("Failed to generate image.")
-        except Exception as e:
-            images.append(f"Error: {str(e)}")
-    return images
+def image_to_data_url(image):
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{img_str}"
 
 def inpaint_defect(image, prompt, num_images=1):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
-    
     image_data_url = image_to_data_url(image)
     images = []
 
@@ -81,7 +81,6 @@ def inpaint_defect(image, prompt, num_images=1):
             "num_inference_steps": 100,
             "image_guidance_scale": 1.5
         }
-
         prediction = rep_client.predictions.create(
             version="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
             input=input
@@ -92,15 +91,38 @@ def inpaint_defect(image, prompt, num_images=1):
             response = requests.get(image_url)
             img = Image.open(io.BytesIO(response.content))
             images.append(img)
-        else:
-            images.append(None)
     return images
 
-def image_to_data_url(image):
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    return f"data:image/png;base64,{img_str}"
+def generate_images(prompts):
+    images = []
+    for prompt in prompts:
+        try:
+            prediction = rep_client.predictions.create(
+                version="ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e4",
+                input={"prompt": prompt, "scheduler": "K_EULER"}
+            )
+            prediction.wait()
+            if prediction.status == "succeeded" and prediction.output:
+                images.append(prediction.output[0])
+        except Exception as e:
+            images.append(f"Error: {str(e)}")
+    return images
+
+def process_railway_defects(prompt, number_of_images):
+    variations = generate_variations(prompt, number_of_images)
+    images = generate_images(variations)
+    return images
+
+def download_images_as_zip(images):
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w') as zf:
+        for i, img in enumerate(images):
+            img_buffer = io.BytesIO()
+            img.save(img_buffer, format='PNG')
+            img_buffer.seek(0)
+            zf.writestr(f'image_{i + 1}.png', img_buffer.read())
+    zip_buffer.seek(0)
+    return zip_buffer
 
 # UI creation
 with gr.Blocks() as app:
@@ -110,53 +132,78 @@ with gr.Blocks() as app:
                 prompt_input = gr.Dropdown(choices=predefined_prompts, label="Select a prompt")
                 number_input_dropdown = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                 submit_button_dropdown = gr.Button("Generate")
-                feedback_input = gr.Textbox(label="Enter your feedback", placeholder="Write your feedback here...")
-                like_button = gr.Button(value="👍 Like", visible=True)
-                dislike_button = gr.Button(value="👎 Dislike", visible=True)
             image_outputs_dropdown = gr.Gallery()
+
             submit_button_dropdown.click(
-                fn=lambda prompt, num: generate_images(generate_variations(prompt, num)),
+                fn=process_railway_defects,
                 inputs=[prompt_input, number_input_dropdown],
-                outputs=[image_outputs_dropdown, feedback_input, like_button, dislike_button]
+                outputs=image_outputs_dropdown
             )
 
         with gr.Tab("Custom Defect"):
             with gr.Row():
                 custom_prompt_input = gr.Textbox(label="Custom Defect")
-                number_input_custom = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum= 10)
+                number_input_custom = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                 submit_button_custom = gr.Button("Generate")
-                feedback_input_custom = gr.Textbox(label="Enter your feedback", placeholder="Write your feedback here...", visible=True)
-                like_button_custom = gr.Button(value="👍 Like", visible=True)
-                dislike_button_custom = gr.Button(value="👎 Dislike", visible=True)
             image_outputs_custom = gr.Gallery()
+
             submit_button_custom.click(
-                fn=lambda prompt, num: generate_images(generate_variations(prompt, num)),
+                fn=process_railway_defects,
                 inputs=[custom_prompt_input, number_input_custom],
-                outputs=[image_outputs_custom, feedback_input_custom, like_button_custom, dislike_button_custom]
+                outputs=image_outputs_custom
             )
-
+        
         with gr.Tab("Inpaint Defect"):
             with gr.Row():
                 image_input = gr.Image(label="Upload Image")
                 inpaint_prompt_input = gr.Textbox(label="Defect Description")
-                number_input_inpaint = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum= 10)
+                number_input_inpaint = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                 submit_button_inpaint = gr.Button("Inpaint Defect")
-                feedback_input_inpaint = gr.Textbox(label="Enter your feedback", placeholder="Write your feedback here...", visible=True)
-                like_button_inpaint = gr.Button(value="👍 Like", visible=True)
-                dislike_button_inpaint = gr.Button(value="👎 Dislike", visible=True)
             inpainted_image_output = gr.Gallery()
+            download_button = gr.Button("Download Images as Zip")
+
             submit_button_inpaint.click(
-                fn=lambda img, prompt, num: inpaint_defect(img, prompt, num),
+                fn=inpaint_defect,
                 inputs=[image_input, inpaint_prompt_input, number_input_inpaint],
-                outputs=[inpainted_image_output, feedback_input_inpaint, like_button_inpaint, dislike_button_inpaint]
+                outputs=inpainted_image_output
             )
 
-    like_button.click(lambda x: handle_feedback(x, "like"), inputs=feedback_input, outputs=None)
-    dislike_button.click(lambda x: handle_feedback(x, "dislike"), inputs=feedback_input, outputs=None)
-    like_button_custom.click(lambda x: handle_feedback(x, "like"), inputs=feedback_input_custom, outputs=None)
-    dislike_button_custom.click(lambda x: handle_feedback(x, "dislike"), inputs=feedback_input_custom, outputs=None)
-    like_button_inpaint.click(lambda x: handle_feedback(x, "like"), inputs=feedback_input_inpaint, outputs=None)
-    dislike_button_inpaint.click(lambda x: handle_feedback(x, "dislike"), inputs=feedback_input_inpaint, outputs=None)
+            download_button.click(
+                fn=download_images_as_zip,
+                inputs=inpainted_image_output,
+                outputs=gr.File(label="Download Zip")
+            )
+        
+        with gr.Tab("Feedback"):
+            feedback_input = gr.Textbox(label="Your feedback")
+            like_button = gr.Button("Like")
+            dislike_button = gr.Button("Dislike")
+            feedback_response = gr.Text(label="Feedback Response")
+
+            def handle_feedback(feedback):
+                sentiment_response = ask_rail_defect_question(feedback, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ')
+                if "negative" in sentiment_response:
+                    return "Sorry for your experience, and thanks for your feedback. We are updating our system to generate more satisfactory samples."
+                else:
+                    return "Thanks for your positive feedback!"
+
+            feedback_input.change(
+                fn=handle_feedback,
+                inputs=feedback_input,
+                outputs=feedback_response
+            )
+
+            like_button.click(
+                fn=lambda: "Thanks for your positive feedback!",
+                inputs=None,
+                outputs=feedback_response
+            )
+
+            dislike_button.click(
+                fn=lambda: "Sorry for your experience, and thanks for your feedback. We are updating our system to generate more satisfactory samples.",
+                inputs=None,
+                outputs=feedback_response
+            )
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch()