Skip to content

Commit

Permalink
fixed issue 32
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexandrequ
alexandrequ committed Mar 25, 2024
1 parent 91d8487 commit 040cae8
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 50 deletions.
8 changes: 7 additions & 1 deletion arcjetCV/utils/processor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,13 @@ def process_all(

# Initialize video writer if write_video is True
if write_video:
video.get_writer()

video_output_name = "video_out_%s_%i_%i.m4v" % (
output_prefix,
first_frame,
last_frame,
)
video.get_writer(video_output_name)

# Setup output JSON file
if output_prefix == "":
Expand Down
121 changes: 72 additions & 49 deletions arcjetCV/utils/video.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,14 @@
import numpy as np
import threading
from arcjetCV.utils.utils import splitfn
from arcjetCV.segmentation.time.time_segmentation import time_segmentation, extract_interest
from arcjetCV.segmentation.time.time_segmentation import (
time_segmentation,
extract_interest,
)


class Video(object):
'''
"""
Convenience wrapper for opencv video capture.
This class provides a simple interface for working with video files using OpenCV. It encapsulates functionality for reading frames from a video file, setting the current frame, and writing processed frames to a new video file.
Expand All @@ -32,7 +35,8 @@ class Video(object):
print(video)
frame = video.get_frame(0)
```
'''
"""

def __init__(self, path):
"""
Initializes the Video object.
Expand All @@ -55,20 +59,27 @@ def __init__(self, path):
self.shape = np.shape(frame)
self.h, self.w, self.chan = self.shape
self.last_frame = frame
print(f'Loaded {self.fpath} with {self.nframes} frames and shape {self.shape}')
print(f"Loaded {self.fpath} with {self.nframes} frames and shape {self.shape}")

if self.chan not in [1, 3]:
raise IndexError("ERROR: number of channels of input video (%i) is not 1 or 3!" % self.chan)
raise IndexError(
"ERROR: number of channels of input video (%i) is not 1 or 3!"
% self.chan
)

### video output
self.writer = None
self.output_path = os.path.join(self.folder,"video_out_"+self.name+'.m4v')
self.output_path = os.path.join(self.folder, "video_out_" + self.name + ".m4v")

def __str__(self):
"""
Returns a string representation of the Video object.
"""
return 'Video: {}, shape={}, nframes={}'.format(self.fpath,self.shape,self.nframes,)
return "Video: {}, shape={}, nframes={}".format(
self.fpath,
self.shape,
self.nframes,
)

def get_next_frame(self):
"""
Expand All @@ -87,7 +98,7 @@ def set_frame(self, index):
:param index: index of the frame to set
"""
with self._lock:
self.cap.set(cv.CAP_PROP_POS_FRAMES,index)
self.cap.set(cv.CAP_PROP_POS_FRAMES, index)
return

def get_frame(self, index):
Expand All @@ -98,7 +109,7 @@ def get_frame(self, index):
:returns: RGB frame at the specified index
"""
with self._lock:
self.cap.set(cv.CAP_PROP_POS_FRAMES,index)
self.cap.set(cv.CAP_PROP_POS_FRAMES, index)
_, self.last_frame = self.cap.read()
return cv.cvtColor(self.last_frame, cv.COLOR_BGR2RGB)

Expand All @@ -110,22 +121,26 @@ def close(self):
self.writer.release()
self.cap.release()

def get_writer(self):
def get_writer(self, video_output_name):
"""
Initializes the video writer.
"""
vid_cod = cv.VideoWriter_fourcc('m','p','4','v')
print(f"Writing {self.output_path}")
self.writer = cv.VideoWriter(self.output_path, vid_cod, self.fps,(self.shape[1], self.shape[0]))
video_output_path = os.path.join(self.folder, video_output_name)
vid_cod = cv.VideoWriter_fourcc("m", "p", "4", "v")
print(f"Writing {video_output_path}")
self.writer = cv.VideoWriter(
video_output_path, vid_cod, self.fps, (self.shape[1], self.shape[0])
)

def close_writer(self):
"""
Closes the video writer.
"""
self.writer.release()


class VideoMeta(dict):
'''
"""
Subclass of dictionary designed to save/load video metadata in JSON format.
This class extends the dictionary class to provide functionality for saving and loading video metadata in JSON format. It also includes methods for resetting frame crop parameters and setting frame crop parameters.
Expand All @@ -142,7 +157,8 @@ class VideoMeta(dict):
video_meta = VideoMeta(video, 'metadata.json') # create/load metadata obj
video_meta.write() # write metadata to file
```
'''
"""

def __init__(self, video, path):
"""
Initializes the VideoMeta object.
Expand All @@ -158,76 +174,80 @@ def __init__(self, video, path):
self.path = path

### Meta Parameters for each video
self['WIDTH'] = None
self['HEIGHT'] = None
self['CHANNELS'] = None
self['NFRAMES'] = None
self['FIRST_GOOD_FRAME'] = None
self['LAST_GOOD_FRAME'] = None
self['MODELPERCENT'] = None
self['FLOW_DIRECTION'] = None
self['PREAMBLE_RANGE'] = None
self['STING_VISIBLE_RANGES'] = None
self['CROP_YMIN'] = None
self['CROP_YMAX'] = None
self['CROP_XMIN'] = None
self['CROP_XMAX'] = None
self['NOTES'] = None
self['BRIGHTNESS'] = None
self["WIDTH"] = None
self["HEIGHT"] = None
self["CHANNELS"] = None
self["NFRAMES"] = None
self["FIRST_GOOD_FRAME"] = None
self["LAST_GOOD_FRAME"] = None
self["MODELPERCENT"] = None
self["FLOW_DIRECTION"] = None
self["PREAMBLE_RANGE"] = None
self["STING_VISIBLE_RANGES"] = None
self["CROP_YMIN"] = None
self["CROP_YMAX"] = None
self["CROP_XMIN"] = None
self["CROP_XMAX"] = None
self["NOTES"] = None
self["BRIGHTNESS"] = None

if os.path.exists(path):
self.load(path)
else:
self["WIDTH"] = video.w
self["HEIGHT"] = video.h
self["CHANNELS"] = video.chan
self["NFRAMES"] = video.nframes

try: # Infer meta parameters
print("Inferring first and last frames ... ", end='')
print("Inferring first and last frames ... ", end="")
_, out = time_segmentation(video)
start, end = extract_interest(out)
print("Done")
self["FIRST_GOOD_FRAME"] = max(round(start[0] * video.nframes / 500), int(video.nframes * 0.1))
self["LAST_GOOD_FRAME"] = min(round(video.nframes * end[-1] / 500), int(video.nframes))
self["FIRST_GOOD_FRAME"] = max(
round(start[0] * video.nframes / 500), int(video.nframes * 0.1)
)
self["LAST_GOOD_FRAME"] = min(
round(video.nframes * end[-1] / 500), int(video.nframes)
)
except:
print("Time Segmentation Failed")
self["FIRST_GOOD_FRAME"] = 0
self["LAST_GOOD_FRAME"] = video.nframes

# initial crop
self.reset_frame_crop()
print("Calculating brightness ... ", end='')
self['BRIGHTNESS'] = []

print("Calculating brightness ... ", end="")
self["BRIGHTNESS"] = []
video.set_frame(0)
for _ in range(video.nframes):
ret, frame = video.cap.read()
if not ret:
break
gray_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
brightness = np.mean(gray_frame)
self['BRIGHTNESS'].append(round(brightness, 2))
self["BRIGHTNESS"].append(round(brightness, 2))
print("Done")
self.write()

def write(self):
"""
Writes the metadata to a JSON file.
"""
print(f"Writing {self.path} file ... ", end='')
fout = open(self.path,'w+')
print(f"Writing {self.path} file ... ", end="")
fout = open(self.path, "w+")
json.dump(self, fout)
fout.close()
print("Done")

def load(self,path):
def load(self, path):
"""
Loads metadata from a JSON file.
:param path: path to the JSON file
"""
fin = open(path,'r')
fin = open(path, "r")
dload = json.load(fin)
self.update(dload)
fin.close()
Expand All @@ -239,7 +259,7 @@ def load(self,path):
self.name = name
self.ext = ext
self.path = path

def reset_frame_crop(self):
"""
Resets frame crop parameters to defaults.
Expand All @@ -248,7 +268,7 @@ def reset_frame_crop(self):
self["CROP_YMAX"] = int(self["HEIGHT"] * 0.90)
self["CROP_XMIN"] = int(self["WIDTH"] * 0.10)
self["CROP_XMAX"] = int(self["WIDTH"] * 0.90)

def set_frame_crop(self, ymin, ymax, xmin, xmax):
"""
Sets frame crop parameters.
Expand All @@ -267,4 +287,7 @@ def crop_range(self):
"""
Returns the crop range.
"""
return [[self['CROP_YMIN'],self['CROP_YMAX']],[self['CROP_XMIN'], self['CROP_XMAX']]]
return [
[self["CROP_YMIN"], self["CROP_YMAX"]],
[self["CROP_XMIN"], self["CROP_XMAX"]],
]

0 comments on commit 040cae8

Please sign in to comment.