import streamlit as st import cv2 import os import numpy as np import torch import torchvision.transforms as T from PIL import Image import io # Ensure this import is present from model.u2net import U2NET # Replace with the actual import path MAX_FILE_SIZE = 5 * 1024 * 1024 # 5MB # Device configuration device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load face detection model face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml') # Preprocessing function def preprocess_image(image): transform = T.Compose([ T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) return transform(image).unsqueeze(0).to(device) # Model loading function def load_model(model, model_path, device): if not os.path.exists(model_path): raise FileNotFoundError(f"Model weights file not found: {model_path}") model.load_state_dict(torch.load(model_path, map_location=device)) model = model.to(device) return model # Initialize U2Net model u2net = U2NET(in_ch=3, out_ch=1) # Load pre-trained model (replace with your model path) try: u2net = load_model(u2net, "u2net_portrait.pth", device) except FileNotFoundError as e: st.error(f"Error: {e}") st.stop() # Function to detect the largest face in an image def detect_single_face(face_cascade, img): gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) faces = face_cascade.detectMultiScale(gray, 1.1, 4) if len(faces) == 0: st.warning("No face detected; processing the entire image.") return None # Filter to keep the largest face largest_face = max(faces, key=lambda b: b[2] * b[3]) return largest_face # Function to crop and resize face region def crop_face(img, face): if face is None: return img (x, y, w, h) = face height, width = img.shape[0:2] # Define padding lpad = int(float(w) * 0.4) rpad = int(float(w) * 0.4) tpad = int(float(h) * 0.6) bpad = int(float(h) * 0.2) left, right = max(x - lpad, 0), min(x + w + rpad, width) top, bottom = max(y - tpad, 0), min(y + h + bpad, height) im_face = img[top:bottom, left:right] if len(im_face.shape) == 2: im_face = np.repeat(im_face[:, :, np.newaxis], 3, axis=2) # Pad to make image square hf, wf = im_face.shape[0:2] if hf > wf: wfp = (hf - wf) // 2 im_face = np.pad(im_face, ((0, 0), (wfp, wfp), (0, 0)), mode='constant', constant_values=255) elif wf > hf: hfp = (wf - hf) // 2 im_face = np.pad(im_face, ((hfp, hfp), (0, 0), (0, 0)), mode='constant', constant_values=255) im_face = cv2.resize(im_face, (512, 512), interpolation=cv2.INTER_AREA) return im_face # Normalize prediction function def normPRED(d): ma = torch.max(d) mi = torch.min(d) return (d - mi) / (ma - mi) # Inference function def inference(net, input): input = input / np.max(input) tmpImg = np.zeros((input.shape[0], input.shape[1], 3)) tmpImg[:, :, 0] = (input[:, :, 2] - 0.406) / 0.225 tmpImg[:, :, 1] = (input[:, :, 1] - 0.456) / 0.224 tmpImg[:, :, 2] = (input[:, :, 0] - 0.485) / 0.229 tmpImg = tmpImg.transpose((2, 0, 1))[np.newaxis, :, :, :] tmpImg = torch.from_numpy(tmpImg).type(torch.FloatTensor) if torch.cuda.is_available(): tmpImg = tmpImg.cuda() with torch.no_grad(): d1, _, _, _, _, _, _ = net(tmpImg) pred = 1.0 - d1[:, 0, :, :] pred = normPRED(pred) pred = pred.squeeze().cpu().data.numpy() return pred # Convert image to pencil drawing def image_to_pencil_drawing(image, line_size, line_density): img_cv = np.array(image) face = detect_single_face(face_cascade, img_cv) im_face = crop_face(img_cv, face) preprocessed_image = preprocess_image(Image.fromarray(im_face)) with torch.no_grad(): output = u2net(preprocessed_image) output_data = output[0].squeeze().cpu().numpy() pencil_drawing = Image.fromarray((output_data * 255).astype(np.uint8), mode='L') img_cv = cv2.cvtColor(np.array(pencil_drawing), cv2.COLOR_GRAY2BGR) gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY) faces = face_cascade.detectMultiScale(gray, 1.1, 4) for (x, y, w, h) in faces: face_roi = img_cv[y:y+h, x:x+w] pencil_drawing = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)) pencil_drawing = pencil_drawing.convert('RGB') pencil_drawing = pencil_drawing.point(lambda p: p * line_density) # Apply Gaussian Blur to adjust line size pencil_drawing_array = np.array(pencil_drawing) blurred = cv2.GaussianBlur(pencil_drawing_array, (line_size, line_size), 0) inverted_image = 255 - blurred inverted_pencil_drawing = Image.fromarray(inverted_image.astype(np.uint8)) return inverted_pencil_drawing # Fix image function to handle default and uploaded images def fix_image(upload=None): if upload: image = Image.open(upload) else: image = Image.open("8.jpg") # Default image path # Create columns for side-by-side display col1, col2 = st.columns(2) # Display the original image with col1: st.image(image, caption="Uploaded Image", use_column_width=True) # Process and display the pencil drawing pencil_drawing = image_to_pencil_drawing(image, line_size, line_density) with col2: st.image(pencil_drawing, caption="Pencil Drawing", use_column_width=True) # Provide download option for the pencil drawing buf = io.BytesIO() pencil_drawing.save(buf, format="PNG") byte_im = buf.getvalue() st.sidebar.download_button( label="Download Pencil Drawing", data=byte_im, file_name="pencil_drawing.png", mime="image/png" ) # Streamlit app setup st.title("Image to Pencil Drawing") # Sidebar for file upload and controls st.sidebar.title("Controls :gear:") uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"]) # Slider controls for line size and density line_size = st.sidebar.slider("Line Size", min_value=1, max_value=15, value=1, step=2) line_density = st.sidebar.slider("Line Density", min_value=0.5, max_value=2.0, value=1.3, step=0.1) # Determine image processing if uploaded_file is not None: if uploaded_file.size > MAX_FILE_SIZE: st.error("The uploaded file is too large. Please upload an image smaller than 5MB.") else: fix_image(upload=uploaded_file) else: fix_image() # Use default image if none uploaded # Add custom CSS for dark theme st.markdown( """ """, unsafe_allow_html=True )