import streamlit as st import os import cv2 import tempfile from PIL import Image from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer, pipeline import torch import pandas as pd from nltk.corpus import wordnet import nltk from transformers import AutoTokenizer, AutoModelForSeq2SeqLM nltk.download('wordnet') nltk.download('omw-1.4') # Load the pre-trained model for image captioning model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09" model = VisionEncoderDecoderModel.from_pretrained(model_name) model.config.pad_token_id feature_extractor = ViTImageProcessor.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) model_sum_name = "google-t5/t5-base" tokenizer_sum = AutoTokenizer.from_pretrained("google-t5/t5-base") model_sum = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base") # Initialize the summarization model summarize_pipe = pipeline("summarization", model=model_sum_name) def generate_caption(image): pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values output_ids = model.generate(pixel_values) caption = tokenizer.decode(output_ids[0], skip_special_tokens=True) return caption def get_synonyms(word): synonyms = set() for syn in wordnet.synsets(word): for lemma in syn.lemmas(): synonyms.add(lemma.name()) return synonyms def search_captions(query, captions): query_words = query.split() query_synonyms = set(query_words) for word in query_words: query_synonyms.update(get_synonyms(word)) results = [] for path, caption in captions.items(): if any(word in caption.split() for word in query_synonyms): results.append((path, caption)) return results def convert_frame_to_pil(frame): return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) def process_video(video_path, frame_interval): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): st.error("Error: Could not open video file.") return [], pd.DataFrame() video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1 frames = [] count = 0 frame_id = 0 while cap.isOpened(): ret, frame = cap.read() if not ret: break if count % frame_interval == 0: frames.append((frame_id, frame)) frame_id += 1 count += 1 if count > video_length - 1: break cap.release() captions_data = [] for i, (frame_id, frame) in enumerate(frames): pil_image = convert_frame_to_pil(frame) caption = generate_caption(pil_image) captions_data.append({'Frame_ID': frame_id + 1, 'Caption': caption}) captions_df = pd.DataFrame(captions_data) return frames, captions_df def image_captioning_page(): st.title("Image Gallery with Captioning and Search") # Sidebar for search functionality with st.sidebar: query = st.text_input("Search images by caption:") # Right side for folder path input and displaying images folder_path = st.text_input("Enter the folder path containing images:") if folder_path and os.path.isdir(folder_path): image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('png', 'jpg', 'jpeg'))] captions = {} for image_file in image_files: image_path = os.path.join(folder_path, image_file) image = Image.open(image_path) caption = generate_caption(image) captions[image_path] = caption # Display images in a 4-column grid cols = st.columns(4) for idx, (image_path, caption) in enumerate(captions.items()): with cols[idx % 4]: st.image(image_path, caption=caption) if query: results = search_captions(query, captions) st.write("Search Results:") for image_path, caption in results: st.image(image_path, caption=caption) # Save captions to CSV if st.button("Save captions to excel"): df = pd.DataFrame(list(captions.items()), columns=['Image', 'Caption']) save_path = st.text_input("Enter the path to save the Excel file:", folder_path) if save_path: if not os.path.exists(save_path): os.makedirs(save_path) excel_file_path = os.path.join(save_path, "captions.xlsx") df.to_excel(excel_file_path, index=False) st.success(f"Captions saved to {excel_file_path}") def live_camera_captioning_page(): st.title("Live Captioning with Webcam") run = st.checkbox('Run') FRAME_WINDOW = st.image([]) camera = cv2.VideoCapture(0) while run: ret, frame = camera.read() if not ret: st.write("Failed to capture image.") break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) FRAME_WINDOW.image(frame) pil_image = Image.fromarray(frame) caption = generate_caption(pil_image) st.write("Caption: ", caption) cv2.waitKey(500) # Capture an image every 0.5 seconds camera.release() def video_captioning_page(): st.title("Video Captioning") # Sidebar for search functionality with st.sidebar: query = st.text_input("Search videos by caption:") # Right side for folder path input and displaying videos folder_path = st.text_input("Enter the folder path containing videos:") if folder_path and os.path.isdir(folder_path): video_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('mp4', 'avi', 'mov', 'mkv'))] captions = {} for video_file in video_files: video_path = os.path.join(folder_path, video_file) frames, captions_df = process_video(video_path, frame_interval=20) if frames and not captions_df.empty: generated_captions = ' '.join(captions_df['Caption']) summary = summarize_pipe(generated_captions)[0]['summary_text'] captions[video_path] = summary # Display videos in a 4-column grid cols = st.columns(4) for idx, (video_path, summary) in enumerate(captions.items()): with cols[idx % 4]: st.video(video_path, caption=summary) if query: results = search_captions(query, captions) st.write("Search Results:") for video_path, summary in results: st.video(video_path, caption=summary) # Save captions to CSV if st.button("Save captions to excel"): df = pd.DataFrame(list(captions.items()), columns=['Video', 'Caption']) save_path = st.text_input("Enter the path to save the Excel file:", folder_path) if save_path: if not os.path.exists(save_path): os.makedirs(save_path) excel_file_path = os.path.join(save_path, "captions.xlsx") df.to_excel(excel_file_path, index=False) st.success(f"Captions saved to {excel_file_path}") def main(): st.sidebar.title("Navigation") page = st.sidebar.selectbox("Select a page", ["Image Captioning", "Live Camera Captioning", "Video Captioning"]) if page == "Image Captioning": image_captioning_page() elif page == "Live Camera Captioning": live_camera_captioning_page() elif page == "Video Captioning": video_captioning_page() if __name__ == "__main__": main()