Update app.py

app.py

@@ -16,10 +16,9 @@ os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
 rep_client = replicate.Client()
 
 # Set your OpenAI API key
-OPENAI_API_KEY = "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"
+OPENAI_API_KEY = "sk-proj-5iy4bwrqAW8GpguiEawaT3BlbkFJ8p88lLSjOCeDbxWsAOlr"
 openai.api_key = OPENAI_API_KEY
 
-# Predefined prompts for the dropdown
 predefined_prompts = [
     "Missing bolts on railway track",
     "Cracks on railway track",
@@ -29,23 +28,12 @@ predefined_prompts = [
     "Damaged railway platform"
 ]
 
-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 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."
 
-# Function to generate variations enhanced by the GPT model
 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"]
     sizes = ["small", "medium", "large", "tiny", "huge"]
@@ -57,61 +45,10 @@ def generate_variations(base_prompt, number_of_variations):
         size = random.choice(sizes)
         weather = random.choice(weather_conditions)
         
-        # Enhance the base prompt with the GPT model
-        enhanced_prompt = base_prompt
-        full_prompt = f"{enhanced_prompt}, with a {size} defect {location}, observed {weather}."
+        full_prompt = f"{base_prompt}, with a {size} defect {location}, observed {weather}."
         variations.append(full_prompt)
     return variations
 
-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}"
-
-# Function to inpaint images
-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 = []
-
-    for _ in range(num_images):
-        input_data = {
-            "input_image": image_data_url,
-            "instruction_text": prompt,
-            "scheduler": "K_EULER_ANCESTRAL",
-            "num_outputs": 1,
-            "guidance_scale": 7.5,
-            "num_inference_steps": 100,
-            "image_guidance_scale": 1.5
-        }
-
-        prediction = rep_client.predictions.create(
-            version="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
-            input=input_data
-        )
-        prediction.wait()
-
-        if prediction.status == "succeeded":
-            image_url = prediction.output[0]
-            print(f"Generated Image URL: {image_url}")  # Debugging line to check the URL
-            try:
-                response = requests.get(image_url)
-                response.raise_for_status()  # Check for HTTP errors
-                img = Image.open(io.BytesIO(response.content))
-                images.append(img)
-            except requests.exceptions.RequestException as e:
-                print(f"Error fetching image from URL: {e}")
-                images.append(None)
-        else:
-            print(f"Prediction failed: {prediction.status}")
-            images.append(None)
-    
-    return images
-
-# Function to generate images from prompts
 def generate_images(prompts):
     images = []
     for prompt in prompts:
@@ -129,87 +66,42 @@ def generate_images(prompts):
             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):
-            if img is not None:
-                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:
+    feedback_input = gr.Textbox(label="Enter your feedback", placeholder="Write your feedback here...")
+    like_button = gr.Button(value="👍 Like")
+    dislike_button = gr.Button(value="👎 Dislike")
+    feedback_result = gr.Textbox(label="System Response", interactive=False)
+    refresh_button = gr.Button("Refresh Page")
+    
     with gr.Tabs("Prompt Input"):
-        with gr.Tab("Current Defects"):
-            with gr.Row():
-                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")
-            image_outputs_dropdown = gr.Gallery()
-
-            def on_submit_click_dropdown(prompt, number_of_images):
-                images = process_railway_defects(prompt, number_of_images)
-                return images
-
-            submit_button_dropdown.click(
-                fn=on_submit_click_dropdown,
-                inputs=[prompt_input, number_input_dropdown],
-                outputs=image_outputs_dropdown
+        with gr.Tab("Generate Images"):
+            prompt_input = gr.Dropdown(choices=predefined_prompts, label="Select a defect prompt")
+            number_input = gr.Number(label="Number of images", value=1, minimum=1, maximum=10)
+            generate_button = gr.Button("Generate")
+            gallery = gr.Gallery(label="Generated Images")
+            
+            generate_button.click(
+                fn=lambda prompt, num: generate_images(generate_variations(prompt, num)),
+                inputs=[prompt_input, number_input],
+                outputs=gallery
             )
 
         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)
-                submit_button_custom = gr.Button("Generate")
+            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)
+            submit_button_custom = gr.Button("Generate")
             image_outputs_custom = gr.Gallery()
 
-            def on_submit_click_custom(custom_prompt, number_of_images):
-                images = process_railway_defects(custom_prompt, number_of_images)
-                return images
-
             submit_button_custom.click(
-                fn=on_submit_click_custom,
+                fn=lambda prompt, num: generate_images(generate_variations(prompt, num)),
                 inputs=[custom_prompt_input, number_input_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)
-                submit_button_inpaint = gr.Button("Inpaint Defect")
-            inpainted_image_output = gr.Gallery()
-            download_button = gr.Button("Download Images as Zip")
 
-            def on_submit_click_inpaint(image, inpaint_prompt, number_of_images):
-                inpainted_images = inpaint_defect(image, inpaint_prompt, num_images=number_of_images)
-                return inpainted_images
-
-            def on_download_click(images):
-                zip_buffer = download_images_as_zip(images)
-                return zip_buffer
-
-            submit_button_inpaint.click(
-                fn=on_submit_click_inpaint,
-                inputs=[image_input, inpaint_prompt_input, number_input_inpaint],
-                outputs=inpainted_image_output
-            )
-
-            download_button.click(
-                fn=on_download_click,
-                inputs=inpainted_image_output,
-                outputs=gr.File(label="Download Zip")
-            )
+    like_button.click(lambda x: handle_feedback(x, "like"), inputs=feedback_input, outputs=feedback_result)
+    dislike_button.click(lambda x: handle_feedback(x, "dislike"), inputs=feedback_input, outputs=feedback_result)
+    refresh_button.click(lambda: gr.update(reload_browser=True))
 
-app.launch()
+if __name__ == "__main__":
+    app.launch()