Create app.py

app.py

@@ -0,0 +1,146 @@
+import gradio as gr
+import replicate
+import os
+import random
+import openai
+import numpy as np
+from PIL import Image
+import requests
+import io
+import base64
+import zipfile
+
+# Set API tokens
+os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
+# Initialize the Replicate client
+rep_client = replicate.Client()
+
+# Set your OpenAI API key
+OPENAI_API_KEY = "sk-proj-5iy4bwrqAW8GpguiEawaT3BlbkFJ8p88lLSjOCeDbxWsAOlr"
+openai.api_key = OPENAI_API_KEY
+
+predefined_prompts = [
+    "Missing bolts on railway track",
+    "Cracks on railway track",
+    "Overgrown vegetation near railway track",
+    "Broken railings on railway bridge",
+    "Debris on railway track",
+    "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 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"]
+    weather_conditions = ["under cold conditions", "during hot weather", "in dry weather", "in humid conditions", "under varying temperatures"]
+
+    variations = []
+    for _ in range(number_of_variations):
+        location = random.choice(locations)
+        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)
+    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 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 = {
+            "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
+        )
+        prediction.wait()
+        if prediction.status == "succeeded":
+            image_url = prediction.output[0]
+            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}"
+
+# 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("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("Inpaint Defect"):
+            image_input = gr.Image(label="Upload Image")
+            inpaint_prompt_input = gr.Textbox(label="Describe the defect")
+            number_input_inpaint = gr.Number(label="Number of images", value=1, minimum=1, maximum=10)
+            inpaint_button = gr.Button("Inpaint Defect")
+            gallery_inpaint = gr.Gallery(label="Inpainted Images")
+
+            inpaint_button.click(
+                fn=lambda img, prompt, num: inpaint_defect(img, prompt, num),
+                inputs=[image_input, inpaint_prompt_input, number_input_inpaint],
+                outputs=gallery_inpaint
+            )
+
+    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))
+
+if __name__ == "__main__":
+    app.launch()