Update app.py

app.py

@@ -1,20 +1,201 @@
+from PIL import Image, ImageEnhance
 
+from transformers import CLIPProcessor, CLIPModel
 
-import gradio as gr
-from transformers import pipeline
+from io import BytesIO
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+# Set OpenAI API Key
+import os
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
-def predict(input_img):
-    predictions = pipeline(input_img)
-    return input_img, {p["label"]: p["score"] for p in predictions} 
+# Load CLIP model and processor
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
 
-gradio_app = gr.Interface(
-    predict,
-    inputs=gr.Image(label="Select hot dog candidate", sources=['upload', 'webcam'], type="pil"),
-    outputs=[gr.Image(label="Processed Image"), gr.Label(label="Result", num_top_classes=2)],
-    title="Hot Dog? Or Not?",
-)
+# Expanded object labels
+object_labels = [
+    "cat", "dog", "house", "tree", "car", "mountain", "flower", "bird", "person", "robot",
+    "a digital artwork", "a portrait", "a landscape", "a futuristic cityscape", "horse", 
+    "lion", "tiger", "elephant", "giraffe", "airplane", "train", "ship", "book", "laptop",
+    "keyboard", "pen", "clock", "cup", "bottle", "backpack", "chair", "table", "sofa", 
+    "bed", "building", "street", "forest", "desert", "waterfall", "sunset", "beach", 
+    "bridge", "castle", "statue", "3D model"
+]
 
-if __name__ == "__main__":
-    gradio_app.launch()
+# Example image for contrast check
+EXAMPLE_IMAGE_URL = "https://www.watercoloraffair.com/wp-content/uploads/2023/04/monet-houses-of-parliament-low-key.jpg"  # Square example image
+example_image = Image.open(BytesIO(requests.get(EXAMPLE_IMAGE_URL).content))
+
+# Function to handle chatbot queries with streaming
+def process_chat(user_text):
+    if not user_text.strip():
+        yield "⚠️ Please enter a question."
+        return
+
+    stream = openai.ChatCompletion.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant named Diane specializing in digital art advice."},
+            {"role": "user", "content": user_text}
+        ],
+        stream=True
+    )
+
+    response = ""
+    for chunk in stream:
+        if chunk.choices[0].delta.get("content"):
+            token = chunk.choices[0].delta["content"]
+            response += token
+            yield response
+
+# Function to analyze image contrast
+def analyze_contrast_opencv(image_path):
+    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+    contrast = img.std()
+    return contrast
+
+# Function to identify objects using CLIP
+def identify_objects_with_clip(image_path):
+    image = Image.open(image_path).convert("RGB")
+    inputs = clip_processor(text=object_labels, images=image, return_tensors="pt", padding=True)
+    with torch.no_grad():
+        outputs = clip_model(**inputs)
+        logits_per_image = outputs.logits_per_image
+        probs = logits_per_image.softmax(dim=1).numpy().flatten()
+    best_match_label = object_labels[probs.argmax()]
+    return best_match_label
+
+# Function to enhance image contrast
+def enhance_contrast(image):
+    enhancer = ImageEnhance.Contrast(image)
+    enhanced_image = enhancer.enhance(1.5)
+    enhanced_path = "enhanced_image.png"
+    enhanced_image.save(enhanced_path)
+    return enhanced_path
+
+# Function to provide additional suggestions with streaming
+def provide_suggestions_streaming(object_identified):
+    if not object_identified:
+        yield "Sorry, I couldn't find an object in your artwork. Try a different image."
+        return
+
+    stream = openai.ChatCompletion.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "You are an expert digital art advisor."},
+            {"role": "user", "content": f"Suggest ways to improve a digital artwork featuring a {object_identified}."}
+        ],
+        stream=True
+    )
+
+    response = ""
+    for chunk in stream:
+        if chunk.choices[0].delta.get("content"):
+            token = chunk.choices[0].delta["content"]
+            response += token
+            yield response
+
+# Main image processing function
+def process_image(image):
+    if not image:
+        return "⚠️ Please upload an image.", None, None
+    image.save("uploaded_image.png")
+    contrast = analyze_contrast_opencv("uploaded_image.png")
+    object_identified = identify_objects_with_clip("uploaded_image.png")
+    if contrast < 25:
+        enhanced_image_path = enhance_contrast(Image.open("uploaded_image.png"))
+        return (
+            f"Hey, great artwork of {object_identified}! However, it looks like the contrast is a little low. I've improved the contrast for you. ✨",
+            enhanced_image_path,
+            object_identified
+        )
+    return (
+        f"Hey, great artwork of {object_identified}! Looks like the color contrast is great. Be proud of yourself! 🌟",
+        None,
+        object_identified
+    )
+
+# Gradio Blocks Interface
+demo = gr.Blocks(css="""
+    #upload-image, #example-image {
+        height: 300px !important;
+    }
+    .button {
+        height: 50px;
+        font-size: 16px;
+    }
+""")
+
+with demo:
+    gr.Markdown("## 🎨 DIANE (Digital Imaging and Art Neural Enhancer)")
+    gr.Markdown("DIANE is here to assist you in refining your digital art. She can answer questions about digital art, analyze your images, and provide creative suggestions to enhance your work.")
+
+    # Chatbot Section
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 💬 Ask me about digital art")
+            user_text = gr.Textbox(label="Enter your question", placeholder="What is the best tool for a beginner?...")
+            chat_output = gr.Textbox(label="Answer", interactive=False)
+            chat_button = gr.Button("Ask", elem_classes="button")
+        
+        chat_button.click(process_chat, inputs=user_text, outputs=chat_output)
+
+    # Image Analysis Section
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 🖼️ Upload an image to check its contrast levels")
+            with gr.Row(equal_height=True):
+                # Left: Image upload field
+                with gr.Column():
+                    image_input = gr.Image(label="Upload an image", type="pil", elem_id="upload-image")
+                    image_button = gr.Button("Check", elem_classes="button")
+                
+                # Right: Example image field
+                with gr.Column():
+                    gr.Image(value=example_image, label="Example Image", interactive=False, elem_id="example-image")
+                    example_button = gr.Button("Use Example Image", elem_classes="button")
+            image_output_text = gr.Textbox(label="Analysis", interactive=False)
+            image_output_image = gr.Image(label="Improved Image", interactive=False)
+            suggestion_button = gr.Button("I want to improve this artwork. Any suggestions?", visible=False)
+            suggestions_output = gr.Textbox(label="Suggestions", interactive=True)
+            state_object = gr.State()  # To store identified object
+
+        # Load example image into the input
+        def use_example_image():
+            return example_image
+        
+        example_button.click(
+            use_example_image,
+            inputs=None,
+            outputs=image_input
+        )
+
+        # Analyze button
+        def update_suggestions_visibility(analysis, enhanced_image, identified_object):
+            return gr.update(visible=True), analysis, enhanced_image
+
+        image_button.click(
+            process_image,
+            inputs=image_input,
+            outputs=[
+                image_output_text, 
+                image_output_image, 
+                state_object
+            ]
+        )
+
+        # Automatically enable suggestions after image processing
+        image_button.click(
+            update_suggestions_visibility,
+            inputs=[image_output_text, image_output_image, state_object],
+            outputs=[suggestion_button, image_output_text, image_output_image]
+        )
+
+        # Suggestion button functionality with streaming
+        suggestion_button.click(
+            provide_suggestions_streaming,
+            inputs=state_object,
+            outputs=suggestions_output
+        )
+
+demo.launch()