first commit

.gitignore

@@ -0,0 +1,3 @@
+*.mp4
+extracted_images/
+modified_images/

.streamlit/config.toml

@@ -0,0 +1,2 @@
+[server]
+maxUploadSize = 50

app.py

@@ -7,6 +7,7 @@ import torch
 import matplotlib.pyplot as plt
 from stqdm import stqdm
 from pathlib import Path
+from moviepy.editor import VideoFileClip
 
 
 # Load the model
@@ -19,6 +20,21 @@ model_pt = AutoModelForObjectDetection.from_pretrained(best_model_path).to(devic
 COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
         [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
 
+
+# Edit video
+def cut_video(clip=None):
+	with st.form("edit"):
+		duration = int(clip.duration)
+		st.write("Edit a small part of video")
+		start = st.sidebar.number_input('Start time (seconds):',max_value=duration)
+		end = st.sidebar.number_input('End time (seconds):',min_value=start+3,max_value=duration)
+		submitted = st.form_submit_button("Edit Out")
+	if submitted:
+		clip = clip.subclip(start, end)
+		clip.write_videofile("edit.mp4")
+		return clip
+
+
 # Convert Video to Frames
 def video_to_frames(video, dir):
     cap = cv2.VideoCapture(str(video))
@@ -35,7 +51,6 @@ def video_to_frames(video, dir):
     cap.release()
     #print (f"No. of frames {frame_count}")
 
-
 # for output bounding box post-processing
 def box_cxcywh_to_xyxy(x):
     x_c, y_c, w, h = x.unbind(1)
@@ -115,12 +130,18 @@ def frames_to_video(dir, path, fps=5):
 		#print(filename)
 		#inserting the frames into an image array
 		frame_array.append(img)
-	out = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*'DIVX'), fps, size)
+	out = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*'mp4v'), fps, size)
 	for item in frame_array:
 		# writing to a image array
 		out.write(item)
 	out.release()
 
+# Display video
+def display(path):
+	video_file = open(str(path), 'rb')
+	video_bytes = video_file.read()
+	st.video(video_bytes)
+
 
 # Main
 if __name__=='__main__':
@@ -140,10 +161,11 @@ if __name__=='__main__':
 		shutil.rmtree(MOD_DIR)
 	MOD_DIR.mkdir(parents=True, exist_ok=True)
 
-	generated_video = Path(BASE_DIR, "final_video.mp4")
+	edited_video = Path(BASE_DIR, "edit.mp4")
+	generated_video = Path(BASE_DIR, "balloons.mp4")
 
 	# Upload the video
-	uploaded_file = st.file_uploader("Upload a small video containing Balloons", type=["mp4"])
+	uploaded_file = st.sidebar.file_uploader("Upload a small video containing Balloons", type=["mp4", "mpeg"])
 	if uploaded_file is not None:
 		st.video(uploaded_file)
 		vid = uploaded_file.name
@@ -152,19 +174,25 @@ if __name__=='__main__':
 			f.write(uploaded_file.read())
 		uploaded_video = Path(BASE_DIR, vid)
 
-		# Detect balloon in the frames and generate video
-		try:
-			video_to_frames(uploaded_video, FRAMES_DIR)
-			predict_on_frames(FRAMES_DIR, MOD_DIR)
-			frames_to_video(MOD_DIR, generated_video)
-			st.success("Successfully Generated!!")
-			
-			# Video file Generated
-			video_file = open(str(generated_video), 'rb')
-			video_bytes = video_file.read()
-			st.video(video_bytes)
-			st.download_button('Download the Video', video_bytes, file_name=generated_video.name)
-		except Exception as e:
-			st.error(f"Could not convert the file due to {e}")
+		clip = VideoFileClip(vid)
+		clip = cut_video(clip)			
+
+		if clip is not None:
+			# Detect balloon in the frames and generate video
+			try:
+				st.info('View Edited Clip')
+				display(edited_video)
+				video_to_frames(edited_video, FRAMES_DIR)
+				predict_on_frames(FRAMES_DIR, MOD_DIR)
+				frames_to_video(MOD_DIR, generated_video)
+				st.success("Successfully Generated!!")
+
+				# Video file Generated
+				display(generated_video)
+				st.download_button('Download the Video', open(str(generated_video), 'rb').read(), file_name=generated_video.name)
+			except Exception as e:
+				st.error(f"Could not convert the file due to {e}")
+		else:
+			st.error("Please submit an edited clip.")
 	else:
 		st.info('File Not Uploaded Yet!!!')

requirements.txt

@@ -5,3 +5,4 @@ matplotlib==3.6.1
 torch==1.12.1
 transformers==4.22.2
 stqdm==0.0.4
+moviepy==1.0.3