Skip to content

Commit a68d948

Browse files
rgeirhosrgeirhos
authored
Merge pull request google-deepmind#15 from sam-motamed/patch-2
Bug fixes + more flexible user interaction
2 parents db94f76 + 62cdb4b commit a68d948

6 files changed

Lines changed: 252 additions & 248 deletions

File tree

‎README.md‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,14 @@ The best possible score on Physics-IQ is 100.0%, this score would be achieved by
3535

3636
| **#** | **Model** | **input type** | **Physics-IQ score** |
3737
| -- | --- | --- | --- |
38-
| 1 | [VideoPoet](https://arxiv.org/abs/2312.14125) | multiframe | **24.1 %** :1st_place_medal: |
39-
| 2 | [Runway Gen 3](https://runwayml.com/research/introducing-gen-3-alpha) | i2v | 18.4 % :2nd_place_medal: |
40-
| 3 | [Lumiere](https://arxiv.org/abs/2401.12945) | multiframe | 18.2 % :3rd_place_medal: |
41-
| 4 | [VideoPoet](https://arxiv.org/abs/2312.14125) | i2v | 18.0 % |
42-
| 5 | [Lumiere](https://arxiv.org/abs/2401.12945) | i2v | 17.1 % |
43-
| 6 | [Stable Video Diffusion](https://arxiv.org/abs/2311.15127) | i2v | 13.5 % |
44-
| 7 | [Pika](https://pika.art/) | i2v | 9.5 % |
45-
| 8 | [Sora](https://openai.com/sora/) | i2v | 8.7 % |
38+
| 1 | [VideoPoet](https://arxiv.org/abs/2312.14125) | multiframe | **29.5 %** :1st_place_medal: |
39+
| 2 | [Lumiere](https://arxiv.org/abs/2401.12945) | multiframe | 23.0 % :2nd_place_medal: |
40+
| 3 | [Runway Gen 3](https://runwayml.com/research/introducing-gen-3-alpha) | i2v | 22.8 % :3rd_place_medal: |
41+
| 4 | [VideoPoet](https://arxiv.org/abs/2312.14125) | i2v | 20.3 % |
42+
| 5 | [Lumiere](https://arxiv.org/abs/2401.12945) | i2v | 19.0 % |
43+
| 6 | [Stable Video Diffusion](https://arxiv.org/abs/2311.15127) | i2v | 14.8 % |
44+
| 7 | [Pika](https://pika.art/) | i2v | 13.0 % |
45+
| 8 | [Sora](https://openai.com/sora/) | i2v | 10.0 % |
4646

4747
---
4848

@@ -160,6 +160,7 @@ python3 code/run_physics_iq.py --input_folders <generated_videos_dirs> --output_
160160

161161
---
162162

163+
163164
## Citation
164165

165166
```latex
@@ -192,4 +193,3 @@ either express or implied. See the licenses for the specific language governing
192193
permissions and limitations under those licenses.
193194

194195
This is not an official Google product.
195-

‎assets/physics_IQ_score.png‎

-101 KB
Binary file not shown.

‎code/binary_mask_generator.py‎

Lines changed: 36 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,25 +19,26 @@
1919
import cv2
2020
import numpy as np
2121

22-
2322
def generate_mask(
24-
in_path: str, out_video_path: str, threshold_value: int = 10
23+
in_path: str, out_video_path: str, is_real: bool, threshold_value: int = 10
2524
) -> None:
2625
"""
2726
Processes a video to generate a binary mask for ROI detection.
2827
2928
Args:
3029
in_path: Path to the input video file.
3130
out_video_path: Path to save the generated binary mask video.
31+
is_real: Boolean indicating whether the video is real.
3232
threshold_value: Threshold value for binary segmentation.
3333
"""
34+
3435
try:
3536
cap = cv2.VideoCapture(in_path)
3637
if not cap.isOpened():
3738
print(f"Error: Could not open video file {in_path}")
3839
return
3940

40-
fps = cap.get(cv2.CAP_PROP_FPS)
41+
fps = round(cap.get(cv2.CAP_PROP_FPS))
4142
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
4243
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
4344
input_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -46,39 +47,45 @@ def generate_mask(
4647
width, height = width - width % 2, height - height % 2
4748

4849
# Use H.264 codec for better compatibility
49-
fourcc = cv2.VideoWriter_fourcc(*"avc1")
50+
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
5051
name_parts = out_video_path.split('/')
5152
name_parts_last = name_parts[-1].split('_')
52-
53-
# Construct the new last part
54-
replace_last_part = (
55-
name_parts_last[0] + '_video-masks_' + f'{int(fps)}FPS_' + name_parts_last[1] + '_take-1_'+ name_parts_last[2]
56-
)
57-
58-
# Replace the last part in the list
53+
54+
#Account for name difference between real and generated videos and save the mask accordingly
55+
if is_real:
56+
replace_last_part = (
57+
name_parts_last[0] + '_video-masks_' + f'{int(fps)}FPS_' + name_parts_last[3] + '_' + name_parts_last[4] + '_'+ name_parts_last[5]
58+
)
59+
else:
60+
replace_last_part = (
61+
name_parts_last[0] + '_video-masks_' + f'{int(fps)}FPS_' + name_parts_last[1] + '_take-1_'+ name_parts_last[2]
62+
)
63+
5964
name_parts[-1] = replace_last_part
60-
61-
# Reconstruct the full path
6265
out_video_path = '/'.join(name_parts)
66+
6367
out = cv2.VideoWriter(
6468
out_video_path, fourcc, fps, (width, height), isColor=False
6569
)
66-
70+
71+
# Read and initialize background model with first real frame
6772
ret, prev_frame = cap.read()
6873
if not ret:
6974
print(f"Error: Could not read the first frame from {in_path}")
7075
return
7176

72-
prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
73-
prev_gray = cv2.GaussianBlur(prev_gray, (5, 5), 0)
77+
# Process the first frame immediately
78+
gray_frame = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
79+
gray_frame = cv2.GaussianBlur(gray_frame, (5, 5), 0)
80+
avg_frame = gray_frame.astype("float")
7481

75-
avg_frame = prev_gray.astype("float")
76-
black_frame = np.zeros((height, width), dtype=np.uint8)
77-
out.write(black_frame)
78-
generated_frame_count = 1
82+
# Write the first frame's mask
83+
first_mask = np.zeros_like(gray_frame, dtype=np.uint8)
84+
out.write(first_mask)
85+
generated_frame_count = 1 # Account for the first written frame
7986

8087
kernel = np.ones((5, 5), np.uint8)
81-
88+
8289
while True:
8390
ret, frame = cap.read()
8491
if not ret:
@@ -98,24 +105,24 @@ def generate_mask(
98105

99106
binary_frame = cv2.morphologyEx(binary_frame, cv2.MORPH_OPEN, kernel)
100107
binary_frame = cv2.morphologyEx(binary_frame, cv2.MORPH_CLOSE, kernel)
101-
102-
# Write the processed frame
103108
out.write(binary_frame)
104109
generated_frame_count += 1
105110

106111
cap.release()
107112
out.release()
108113
print(f"Processed video saved at {out_video_path}")
109114
print(f"Input video frame count: {input_frame_count}")
110-
print(f"Number of generated frames: {generated_frame_count}")
115+
print(f"Number of generated mask frames: {generated_frame_count}")
111116

112117
except cv2.error as e:
113118
print(f"Error processing video {in_path}: {e}")
114119

115120

121+
116122
def generate_binary_masks(
117123
input_parent_directory: str,
118124
output_parent_directory: str,
125+
is_real: bool,
119126
threshold_value: int = 10,
120127
) -> None:
121128
"""
@@ -124,8 +131,10 @@ def generate_binary_masks(
124131
Args:
125132
input_parent_directory: Path to the parent directory containing input videos.
126133
output_parent_directory: Path to the parent directory for saving output videos.
134+
is_real: Boolean indicating whether the video is real.
127135
threshold_value: Threshold value for binary segmentation.
128136
"""
137+
129138
if not os.path.exists(output_parent_directory):
130139
os.makedirs(output_parent_directory)
131140

@@ -139,7 +148,7 @@ def generate_binary_masks(
139148
input_path = os.path.join(root, filename)
140149
output_video_path = os.path.join(output_directory, filename)
141150
print(f"Generating mask for: {filename}")
142-
generate_mask(input_path, output_video_path, threshold_value)
151+
generate_mask(input_path, output_video_path, is_real, threshold_value)
143152

144153

145154
if __name__ == "__main__":
@@ -164,11 +173,8 @@ def generate_binary_masks(
164173
default=10,
165174
help="Threshold value for binary segmentation.",
166175
)
167-
168176
args = parser.parse_args()
169-
177+
170178
generate_binary_masks(
171-
args.input_parent_directory,
172-
args.output_parent_directory,
173-
args.threshold_value,
179+
args.input_parent_directory, args.output_parent_directory, True, args.threshold_value
174180
)

0 commit comments

Comments
 (0)