Update app.py

app.py

@@ -1,123 +1,87 @@
 import streamlit as st
-import tensorflow as tf
-import numpy as np
+from fastai.vision import open_image, load_learner, show_image
+import PIL.Image
 from PIL import Image
-import requests
 from io import BytesIO
+import requests
+import torch.nn as nn
+import os
+import tempfile
+import shutil
+
+# Define the FeatureLoss class
+class FeatureLoss(nn.Module):
+    def __init__(self, m_feat, layer_ids, layer_wgts):
+        super().__init__()
+        self.m_feat = m_feat
+        self.loss_features = [self.m_feat[i] for i in layer_ids]
+        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.wgts = layer_wgts
+        self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))] + [f'gram_{i}' for i in range(len(layer_ids))]
+
+    def make_features(self, x, clone=False):
+        self.m_feat(x)
+        return [(o.clone() if clone else o) for o in self.hooks.stored]
+
+    def forward(self, input, target):
+        out_feat = self.make_features(target, clone=True)
+        in_feat = self.make_features(input)
+        self.feat_losses = [base_loss(input, target)]
+        self.feat_losses += [base_loss(f_in, f_out) * w for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
+        self.feat_losses += [base_loss(gram_matrix(f_in), gram_matrix(f_out)) * w**2 * 5e3 for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
+        self.metrics = dict(zip(self.metric_names, self.feat_losses))
+        return sum(self.feat_losses)
+
+    def __del__(self): self.hooks.remove()
+
+def add_margin(pil_img, top, right, bottom, left, color):
+    width, height = pil_img.size
+    new_width = width + right + left
+    new_height = height + top + bottom
+    result = Image.new(pil_img.mode, (new_width, new_height), color)
+    result.paste(pil_img, (left, top))
+    return result
+
+def inference(image_path_or_url, learn):
+    if image_path_or_url.startswith('http://') or image_path_or_url.startswith('https://'):
+        response = requests.get(image_path_or_url)
+        img = PIL.Image.open(BytesIO(response.content)).convert("RGB")
+    else:
+        img = PIL.Image.open(image_path_or_url).convert("RGB")
+
+    im_new = add_margin(img, 250, 250, 250, 250, (255, 255, 255))
+    im_new.save("test.jpg", quality=95)
+    img = open_image("test.jpg")
+    p, img_hr, b = learn.predict(img)
+    return img_hr
+
+# Streamlit application
+st.title("Image Inference with Fastai")
 
 # Download the model file from the Hugging Face repository
-model_url = "https://huggingface.co/Hammad712/Emotion_Detection/resolve/main/emotion_detection_model.h5"
-response = requests.get(model_url)
-model_file_path = 'emotion_detection_model.h5'
-
-# Save the downloaded model file
-with open(model_file_path, 'wb') as f:
-    f.write(response.content)
-
-# Load the saved model
-model = tf.keras.models.load_model(model_file_path)
-
-# Define the emotion labels
-emotion_labels = ['happy', 'sad', 'neutral']
-
-# Function to preprocess the image for the model
-def preprocess_image(image):
-    image = image.resize((48, 48))
-    image = image.convert('L')
-    image = np.array(image)
-    image = image / 255.0
-    image = np.expand_dims(image, axis=-1)
-    image = np.expand_dims(image, axis=0)
-    return image
-
-# Custom CSS for styling
-st.markdown("""
-    <style>
-    .title {
-        font-size: 3em;
-        color: #4CAF50; /* Green */
-        text-align: center;
-    }
-    .description {
-        font-size: 1.2em;
-        color: #777777; /* Gray */
-        text-align: center;
-    }
-    .header {
-        font-size: 1.5em;
-        color: #ff6f61; /* Red */
-    }
-    .predicted-emotion {
-        font-size: 1.5em;
-        color: #1e90ff; /* Blue */
-    }
-    .spinner-text {
-        font-size: 1.2em;
-        color: #ffa500; /* Orange */
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# Title and description
-st.markdown('<div class="title">Emotion Detection</div>', unsafe_allow_html=True)
-st.markdown('<div class="description">Upload an image or provide an image URL to detect emotions</div>', unsafe_allow_html=True)
-
-# Input field for image URL
-st.markdown('<div class="header">Enter Image URL</div>', unsafe_allow_html=True)
-image_url = st.text_input("Image URL")
-
-# File uploader for images
-st.markdown('<div class="header">Or Upload Image</div>', unsafe_allow_html=True)
-image_file = st.file_uploader("Choose an image", type=["png", 'jpg'])
-
-# Sidebar to show previous predictions
-if 'predictions' not in st.session_state:
-    st.session_state.predictions = []
-
-st.sidebar.header("Previous Predictions")
-for pred in st.session_state.predictions:
-    st.sidebar.image(pred['image'], caption=pred['emotion'], use_column_width=True)
-
-# Function to load image from URL
-def load_image_from_url(url):
-    try:
-        response = requests.get(url)
-        img = Image.open(BytesIO(response.content))
-        return img
-    except Exception as e:
-        st.error(f"Error loading image: {e}")
-        return None
-
-# Display the image and predict emotions
-image = None
-if image_url:
-    image = load_image_from_url(image_url)
-elif image_file is not None:
-    image = Image.open(image_file)
-
-if image:
-    st.image(image, caption="Selected Image", use_column_width=True)
-
-    with st.spinner('Processing...'):
-        st.markdown('<div class="spinner-text">Processing...</div>', unsafe_allow_html=True)
-
-        # Preprocess the image
-        preprocessed_image = preprocess_image(image)
-
-        # Make predictions
-        predictions = model.predict(preprocessed_image)
-        predicted_label = emotion_labels[np.argmax(predictions)]
-
-        # Display the predicted emotion
-        st.markdown(f'<div class="predicted-emotion">Predicted Emotion: {predicted_label}</div>', unsafe_allow_html=True)
-
-        # Save the prediction and image to session state
-        image_bytes = BytesIO()
-        image.save(image_bytes, format='PNG')
-        st.session_state.predictions.append({
-            'image': image_bytes.getvalue(),
-            'emotion': predicted_label
-        })
-
-        # Update the sidebar with the new prediction
-        st.sidebar.image(image_bytes.getvalue(), caption=predicted_label, use_column_width=True)
+model_url = "https://huggingface.co/Hammad712/image2sketch/resolve/main/image2sketch.pkl"
+model_file_path = 'image2sketch.pkl'
+
+if not os.path.exists(model_file_path):
+    with st.spinner('Downloading model...'):
+        response = requests.get(model_url)
+        with open(model_file_path, 'wb') as f:
+            f.write(response.content)
+    st.success('Model downloaded successfully!')
+
+# Create a temporary directory for the model
+with tempfile.TemporaryDirectory() as tmpdirname:
+    shutil.move(model_file_path, os.path.join(tmpdirname, 'export.pkl'))
+    learn = load_learner(tmpdirname)
+
+# Input for image URL or path
+image_path_or_url = st.text_input("Enter image path or URL", "")
+
+# Run inference button
+if st.button("Run Inference"):
+    if image_path_or_url:
+        with st.spinner('Processing...'):
+            high_res_image = inference(image_path_or_url, learn)
+            st.image(high_res_image, caption='High Resolution Image', use_column_width=True)
+    else:
+        st.error("Please enter a valid image path or URL.")